import numpy as np
import torch as pt
import torchvision as ptv
from torch.utils.data import DataLoader, Dataset, Subset
from torchvision.datasets import MNIST
import os
import sys
import matplotlib.pyplot as plt

np.random.seed(777)
pt.manual_seed(777)

BATCH_SIZE = 16
DATA_DIR = '../../../../../large_data/DL2/pt/mnist'
mnist_train = ptv.datasets.MNIST(DATA_DIR, train=True,
                                 transform=ptv.transforms.ToTensor(),
                                 download=False)
mnist_train_transf = ptv.datasets.MNIST(DATA_DIR, train=True,
                                        transform=ptv.transforms.Compose([
                                            # ptv.transforms.RandomHorizontalFlip(p=1.0),
                                            # ptv.transforms.RandomVerticalFlip(p=1.0),
                                            # ptv.transforms.RandomRotation(degrees=45),
                                            # ptv.transforms.RandomRotation(degrees=(45, 45)),
                                            ptv.transforms.Resize((32, 32)),
                                            ptv.transforms.RandomCrop((28, 28)),
                                            ptv.transforms.ToTensor(),
                                            ptv.transforms.Normalize((0.), (1.))
                                        ]),
                                        download=False)
dl = DataLoader(mnist_train, batch_size=BATCH_SIZE, shuffle=False)
dl_transf = DataLoader(mnist_train_transf, batch_size=BATCH_SIZE, shuffle=False)

plt.figure(figsize=[12, 6])
spr = 4
spc = 8
spn = 0

for (bx, by), (bx2, by2) in zip(dl, dl_transf):
    for i, bxi in enumerate(bx):
        spn += 1
        if (spn > spr * spc):
            break
        plt.subplot(spr, spc, spn)
        bxi = bxi.transpose(0, 2)
        bxi = bxi.transpose(0, 1)
        plt.imshow(bxi)
        plt.axis('off')
        cls_id = by[i].item()
        plt.title(str(cls_id))

        spn += 1
        if (spn > spr * spc):
            break
        plt.subplot(spr, spc, spn)
        bxi = bx2[i]
        bxi = bxi.transpose(0, 2)
        bxi = bxi.transpose(0, 1)
        plt.imshow(bxi)
        plt.axis('off')
        cls_id = by2[i].item()
        plt.title(str(cls_id))
    if (spn > spr * spc):
        break
